#include "sgd.h"
#include "cmath"

#include "iostream"

/**
 * @brief      return absolute value
 *
 * @param      x     input value
 */
#define sgd_fabs(x) ((x < 0) ? -1*x : x)

/**
 * @brief      return the bigger value
 *
 * @param      a     input value
 * @param      b     another input value
 *
 * @return     the bigger value
 */
#define sgd_max(a, b) (a>b?a:b)

/**
 * @brief      return value of the sgd_solver() function.
 */
enum sgd_return_e
{
	SGD_SUCCESS = 0, ///< The optimization terminated successfully.
	SGD_CONVERGENCE = 1, ///< The optimization reached convergence.
	SGD_STOP, ///< The process stopped by the monitoring function.
	SGD_UNKNOWN_ERROR = -1024, ///< Unknown error.
	SGD_INVALID_VARIABLE_SIZE, ///< The variable size is negative
	SGD_INVALID_MAX_ITERATIONS, ///< The maximal iteration times is negative.
	SGD_INVALID_EPSILON, ///< The epsilon is negative.
	SGD_REACHED_MAX_ITERATIONS, ///< Iteration reached max limit.
	SGD_INVALID_MU, ///< Invalid value for mu.
	SGD_INVALID_ALPHA, ///< Invalid value for alpha.
	SGD_INVALID_BETA, ///< Invalid value for beta.
	SGD_INVALID_SIGMA, ///< Invalid value for sigma.
	SGD_NAN_VALUE, ///< Nan value.
};

/**
 * Default parameter for the SGD methods.
 */
static const sgd_para defparam = {300, 1e-6, 0.01, 0.001, 0.9, 0.999, 1e-8};

sgd_float *sgd_malloc(const int n_size)
{
	sgd_float *x = new sgd_float [n_size];
	return x;
}

void sgd_free(sgd_float* x)
{
	if (x != nullptr) delete[] x;
	x = nullptr;
	return;
}

sgd_para sgd_default_parameters()
{
	sgd_para param = defparam;
	return param;
}

const char* sgd_error_str(int er_index)
{
	switch (er_index)
	{
		case SGD_SUCCESS:
			return "Success.";
		case SGD_CONVERGENCE:
			return "The iteration reached convergence.";
		case SGD_STOP:
			return "The iteration stopped by the progress evaluation function.";
		case SGD_UNKNOWN_ERROR:
			return "Unknown error.";
		case SGD_INVALID_VARIABLE_SIZE:
			return "Invalid array size.";
		case SGD_INVALID_MAX_ITERATIONS:
			return "Invalid maximal iteration times.";
		case SGD_REACHED_MAX_ITERATIONS:
			return "The maximal iteration is reached.";
		case SGD_INVALID_EPSILON:
			return "Invalid value for epsilon.";
		case SGD_INVALID_BETA:
			return "Invalid value for beta.";
		case SGD_INVALID_MU:
			return "Invalid value for mu.";
		case SGD_INVALID_ALPHA:
			return "Invalid value for alpha.";
		case SGD_INVALID_SIGMA:
			return "Invalid value for sigma.";
		case SGD_NAN_VALUE:
			return "NaN values found.";
		default:
			return "Unknown error.";
	}
}

/**
 * @brief      An SGD solver function.
 * 
 * @note       The size of all arrays must be equal to n_size.
 *
 * @param[in]  Evafp       Callback function for calculating the objective function and its gradient.
 * @param[in]  Profp       Callback function for monitoring the optimization process.
 * @param      fx          Returned best value of the objective function by now.
 * @param      m           Pointer of the solution array.
 * @param[in]  n_size      Length of the solution array.
 * @param[in]  m_size      Length of the observation.
 * @param[in]  param       Parameters of optimization process.
 * @param      instance    The user data sent for the function by the client.
 *
 * @return     Status of the function.
 */
typedef int (*sgd_solver_ptr)(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *m, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);

int momentum(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);
int nag(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);
int adagrad(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);
int rmsprop(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);
int adam(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);
int nadam(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);
int adamax(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);
int adabelief(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance);


int sgd_solver(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *m, 
	const int n_size, const int m_size, const sgd_para *param, void *instance, sgd_solver_enum solver_id)
{
	sgd_solver_ptr solver;
	switch (solver_id)
	{
		case SGD_MOMENTUM:
			solver = momentum;
			break;
		case SGD_NAG:
			solver = nag;
			break;
		case SGD_ADAGRAD:
			solver = adagrad;
			break;
		case SGD_RMSPROP:
			solver = rmsprop;
			break;
		case SGD_ADAM:
			solver = adam;
			break;
		case SGD_NADAM:
			solver = nadam;
			break;
		case SGD_ADAMAX:
			solver = adamax;
			break;
		case SGD_ADABELIEF:
			solver = adabelief;
			break;
		default:
			solver = adam;
			break;
	}

	return solver(Evafp, Profp, fx, m, n_size, m_size, param, instance);
}

int momentum(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance)
{
	// set the Adam's parameters
	sgd_para para = (param != nullptr) ? (*param) : defparam;

	//check parameters
	if (n_size <= 0) return SGD_INVALID_VARIABLE_SIZE;
	if (para.iteration <= 0) return SGD_INVALID_MAX_ITERATIONS;
	if (para.epsilon < 0) return SGD_INVALID_EPSILON;
	if (para.mu < 0 || para.mu >= 1.0) return SGD_INVALID_MU;

	sgd_float *mk = sgd_malloc(n_size);
	sgd_float *g  = sgd_malloc(n_size);

	for (int i = 0; i < n_size; i++)
	{
		mk[i] = 0.0;
	}

	sgd_float tmp_fx;
	for (int t = 0; t < para.iteration; t++)
	{
		*fx = 0.0;
		for (int m = 0; m < m_size; m++)
		{
			tmp_fx = Evafp(instance, x, g, n_size, m);
			*fx += tmp_fx;

			for (int i = 0; i < n_size; i++)
			{
				mk[i] = para.mu*mk[i] + g[i];

				x[i] = x[i] - para.alpha * mk[i];
				if (x[i] != x[i]) return SGD_NAN_VALUE;
			}
		}
		*fx /= m_size;

		if (Profp(instance, *fx, x, g, param, n_size, t)) return SGD_STOP;
		if (*fx < para.epsilon) return SGD_CONVERGENCE;
	}

	sgd_free(mk);
	sgd_free(g);
	return SGD_REACHED_MAX_ITERATIONS;
}

int nag(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance)
{
	// set the Adam's parameters
	sgd_para para = (param != nullptr) ? (*param) : defparam;

	//check parameters
	if (n_size <= 0) return SGD_INVALID_VARIABLE_SIZE;
	if (para.iteration <= 0) return SGD_INVALID_MAX_ITERATIONS;
	if (para.epsilon < 0) return SGD_INVALID_EPSILON;
	if (para.mu < 0 || para.mu >= 1.0) return SGD_INVALID_MU;

	sgd_float *mk = sgd_malloc(n_size);
	sgd_float *xk = sgd_malloc(n_size);
	sgd_float *g  = sgd_malloc(n_size);

	for (int i = 0; i < n_size; i++)
	{
		mk[i] = 0.0;
	}

	sgd_float tmp_fx;
	for (int t = 0; t < para.iteration; t++)
	{
		*fx = 0.0;
		for (int m = 0; m < m_size; m++)
		{
			for (int i = 0; i < n_size; i++)
			{
				xk[i] = x[i] - para.mu*para.alpha*mk[i];
			}

			tmp_fx = Evafp(instance, xk, g, n_size, m);
			*fx += tmp_fx;

			for (int i = 0; i < n_size; i++)
			{
				mk[i] = para.mu*mk[i] + g[i];

				x[i] = x[i] - para.alpha * mk[i];
				if (x[i] != x[i]) return SGD_NAN_VALUE;
			}
		}
		*fx /= m_size;

		if (Profp(instance, *fx, x, g, param, n_size, t)) return SGD_STOP;
		if (*fx < para.epsilon) return SGD_CONVERGENCE;
	}

	sgd_free(mk);
	sgd_free(xk);
	sgd_free(g);
	return SGD_REACHED_MAX_ITERATIONS;
}

int adagrad(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance)
{
	// set the Adam's parameters
	sgd_para para = (param != nullptr) ? (*param) : defparam;

	//check parameters
	if (n_size <= 0) return SGD_INVALID_VARIABLE_SIZE;
	if (para.iteration <= 0) return SGD_INVALID_MAX_ITERATIONS;
	if (para.epsilon < 0.0) return SGD_INVALID_EPSILON;
	if (para.sigma < 0.0) return SGD_INVALID_SIGMA;

	sgd_float *mk = sgd_malloc(n_size);
	sgd_float *g  = sgd_malloc(n_size);

	for (int i = 0; i < n_size; i++)
	{
		mk[i] = 0.0;
	}

	sgd_float tmp_fx;
	for (int t = 0; t < para.iteration; t++)
	{
		*fx = 0.0;
		for (int m = 0; m < m_size; m++)
		{
			tmp_fx = Evafp(instance, x, g, n_size, m);
			*fx += tmp_fx;

			for (int i = 0; i < n_size; i++)
			{
				mk[i] = mk[i] + g[i]*g[i];

				x[i] = x[i] - para.alpha * g[i]/(sqrt(mk[i]) + para.sigma);
				if (x[i] != x[i]) return SGD_NAN_VALUE;
			}
		}
		*fx /= m_size;

		if (Profp(instance, *fx, x, g, param, n_size, t)) return SGD_STOP;
		if (*fx < para.epsilon) return SGD_CONVERGENCE;
	}

	sgd_free(mk);
	sgd_free(g);
	return SGD_REACHED_MAX_ITERATIONS;
}

int rmsprop(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance)
{
	// set the Adam's parameters
	sgd_para para = (param != nullptr) ? (*param) : defparam;

	//check parameters
	if (n_size <= 0) return SGD_INVALID_VARIABLE_SIZE;
	if (para.iteration <= 0) return SGD_INVALID_MAX_ITERATIONS;
	if (para.epsilon < 0.0) return SGD_INVALID_EPSILON;
	if (para.sigma < 0.0) return SGD_INVALID_SIGMA;

	sgd_float *vk = sgd_malloc(n_size);
	sgd_float *g  = sgd_malloc(n_size);

	for (int i = 0; i < n_size; i++)
	{
		vk[i] = 0.0;
	}

	sgd_float tmp_fx;
	for (int t = 0; t < para.iteration; t++)
	{
		*fx = 0.0;
		for (int m = 0; m < m_size; m++)
		{
			tmp_fx = Evafp(instance, x, g, n_size, m);
			*fx += tmp_fx;

			for (int i = 0; i < n_size; i++)
			{
				vk[i] = para.beta_2 * vk[i] + (1.0 - para.beta_2)*g[i]*g[i];

				x[i] = x[i] - para.alpha * g[i]/(sqrt(vk[i]) + para.sigma);
				if (x[i] != x[i]) return SGD_NAN_VALUE;
			}
		}
		*fx /= m_size;

		if (Profp(instance, *fx, x, g, param, n_size, t)) return SGD_STOP;
		if (*fx < para.epsilon) return SGD_CONVERGENCE;
	}

	sgd_free(vk);
	sgd_free(g);
	return SGD_REACHED_MAX_ITERATIONS;
}

int adam(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance)
{
	// set the Adam's parameters
	sgd_para para = (param != nullptr) ? (*param) : defparam;

	//check parameters
	if (n_size <= 0) return SGD_INVALID_VARIABLE_SIZE;
	if (para.iteration <= 0) return SGD_INVALID_MAX_ITERATIONS;
	if (para.epsilon < 0) return SGD_INVALID_EPSILON;
	if (para.alpha < 0) return SGD_INVALID_ALPHA;
	if (para.beta_1 < 0.0 || para.beta_1 >= 1.0) return SGD_INVALID_BETA;
	if (para.beta_2 < 0.0 || para.beta_2 >= 1.0) return SGD_INVALID_BETA;
	if (para.sigma < 0.0) return SGD_INVALID_SIGMA;

	sgd_float *mk = sgd_malloc(n_size);
	sgd_float *vk = sgd_malloc(n_size);
	sgd_float *g  = sgd_malloc(n_size);

	for (int i = 0; i < n_size; i++)
	{
		mk[i] = vk[i] = 0.0;
	}

	sgd_float beta_1t = 1.0, beta_2t = 1.0;
	sgd_float alpha_k;

	sgd_float tmp_fx;
	for (int t = 0; t < para.iteration; t++)
	{
		beta_1t *= para.beta_1;
		beta_2t *= para.beta_2;

		alpha_k = para.alpha * sqrt(1.0 - beta_2t)/(1.0 - beta_1t);

		*fx = 0.0;
		for (int m = 0; m < m_size; m++)
		{
			tmp_fx = Evafp(instance, x, g, n_size, m);
			*fx += tmp_fx;

			for (int i = 0; i < n_size; i++)
			{
				mk[i] = para.beta_1*mk[i] + (1.0 - para.beta_1)*g[i];
				vk[i] = para.beta_2*vk[i] + (1.0 - para.beta_2)*g[i]*g[i];

				x[i] = x[i] - alpha_k * mk[i]/(sqrt(vk[i]) + para.sigma);
				if (x[i] != x[i]) return SGD_NAN_VALUE;
			}
		}
		*fx /= m_size;

		if (Profp(instance, *fx, x, g, param, n_size, t)) return SGD_STOP;
		if (*fx < para.epsilon) return SGD_CONVERGENCE;
	}

	sgd_free(mk);
	sgd_free(vk);
	sgd_free(g);
	return SGD_REACHED_MAX_ITERATIONS;
}

int nadam(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance)
{
	// set the Adam's parameters
	sgd_para para = (param != nullptr) ? (*param) : defparam;

	//check parameters
	if (n_size <= 0) return SGD_INVALID_VARIABLE_SIZE;
	if (para.iteration <= 0) return SGD_INVALID_MAX_ITERATIONS;
	if (para.epsilon < 0) return SGD_INVALID_EPSILON;
	if (para.alpha < 0) return SGD_INVALID_ALPHA;
	if (para.beta_1 < 0.0 || para.beta_1 >= 1.0) return SGD_INVALID_BETA;
	if (para.beta_2 < 0.0 || para.beta_2 >= 1.0) return SGD_INVALID_BETA;
	if (para.sigma < 0.0) return SGD_INVALID_SIGMA;

	sgd_float *mk = sgd_malloc(n_size);
	sgd_float *mk_hat = sgd_malloc(n_size);
	sgd_float *nk = sgd_malloc(n_size);
	sgd_float *nk_hat = sgd_malloc(n_size);
	sgd_float *g  = sgd_malloc(n_size);
	sgd_float *g_hat = sgd_malloc(n_size);

	for (int i = 0; i < n_size; i++)
	{
		mk[i] = nk[i] = 0.0;
	}

	sgd_float beta_1t = 1.0, beta_1t1 = para.beta_1, beta_2t = 1.0;

	sgd_float tmp_fx;
	for (int t = 0; t < para.iteration; t++)
	{
		beta_1t  *= para.beta_1;
		beta_1t1 *= para.beta_1;
		beta_2t  *= para.beta_2;

		*fx = 0.0;
		for (int m = 0; m < m_size; m++)
		{
			tmp_fx = Evafp(instance, x, g, n_size, m);
			*fx += tmp_fx;

			for (int i = 0; i < n_size; i++)
			{
				g_hat[i] = g[i]/(1.0 - beta_1t);
				mk[i] = para.beta_1*mk[i] + (1.0 - para.beta_1)*g[i];
				nk[i] = para.beta_2*nk[i] + (1.0 - para.beta_2)*g[i]*g[i];

				mk_hat[i] = mk[i]/(1.0 - beta_1t1);
				nk_hat[i] = nk[i]/(1.0 - beta_2t);

				x[i] = x[i] - para.alpha * ((1.0 - beta_1t)*g_hat[i] 
					+ beta_1t1*mk_hat[i])/(sqrt(nk_hat[i]) + para.sigma);
				if (x[i] != x[i]) return SGD_NAN_VALUE;
			}
		}
		*fx /= m_size;

		if (Profp(instance, *fx, x, g, param, n_size, t)) return SGD_STOP;
		if (*fx < para.epsilon) return SGD_CONVERGENCE;
	}

	sgd_free(mk);
	sgd_free(mk_hat);
	sgd_free(nk);
	sgd_free(nk_hat);
	sgd_free(g);
	sgd_free(g_hat);
	return SGD_REACHED_MAX_ITERATIONS;
}

int adamax(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance)
{
	// set the Adam's parameters
	sgd_para para = (param != nullptr) ? (*param) : defparam;
	if (param == nullptr) para.alpha = 0.002;

	//check parameters
	if (n_size <= 0) return SGD_INVALID_VARIABLE_SIZE;
	if (para.iteration <= 0) return SGD_INVALID_MAX_ITERATIONS;
	if (para.epsilon < 0) return SGD_INVALID_EPSILON;
	if (para.alpha < 0) return SGD_INVALID_ALPHA;
	if (para.beta_1 < 0.0 || para.beta_1 >= 1.0) return SGD_INVALID_BETA;
	if (para.beta_2 < 0.0 || para.beta_2 >= 1.0) return SGD_INVALID_BETA;
	if (para.sigma < 0.0) return SGD_INVALID_SIGMA;

	sgd_float *mk = sgd_malloc(n_size);
	sgd_float *vk = sgd_malloc(n_size);
	sgd_float *g  = sgd_malloc(n_size);

	for (int i = 0; i < n_size; i++)
	{
		mk[i] = vk[i] = 0.0;
	}

	sgd_float beta_1t = 1.0;

	sgd_float tmp_fx;
	for (int t = 0; t < para.iteration; t++)
	{
		beta_1t *= para.beta_1;

		*fx = 0.0;
		for (int m = 0; m < m_size; m++)
		{
			tmp_fx = Evafp(instance, x, g, n_size, m);
			*fx += tmp_fx;

			for (int i = 0; i < n_size; i++)
			{
				mk[i] = para.beta_1*mk[i] + (1.0 - para.beta_1)*g[i];
				vk[i] = sgd_max(para.beta_2*vk[i], sgd_fabs(g[i]));

				x[i] = x[i] - para.alpha * mk[i]/((1.0 - beta_1t)*vk[i]);
				if (x[i] != x[i]) return SGD_NAN_VALUE;
			}
		}
		*fx /= m_size;

		if (Profp(instance, *fx, x, g, param, n_size, t)) return SGD_STOP;
		if (*fx < para.epsilon) return SGD_CONVERGENCE;
	}

	sgd_free(mk);
	sgd_free(vk);
	sgd_free(g);
	return SGD_REACHED_MAX_ITERATIONS;
}

int adabelief(sgd_evaulate_ptr Evafp, sgd_progress_ptr Profp, sgd_float *fx, sgd_float *x, 
	const int n_size, const int m_size, const sgd_para *param, void *instance)
{
	// set the Adam's parameters
	sgd_para para = (param != nullptr) ? (*param) : defparam;

	//check parameters
	if (n_size <= 0) return SGD_INVALID_VARIABLE_SIZE;
	if (para.iteration <= 0) return SGD_INVALID_MAX_ITERATIONS;
	if (para.epsilon < 0) return SGD_INVALID_EPSILON;
	if (para.alpha < 0) return SGD_INVALID_ALPHA;
	if (para.beta_1 < 0.0 || para.beta_1 >= 1.0) return SGD_INVALID_BETA;
	if (para.beta_2 < 0.0 || para.beta_2 >= 1.0) return SGD_INVALID_BETA;
	if (para.sigma < 0.0) return SGD_INVALID_SIGMA;

	sgd_float *mk = sgd_malloc(n_size);
	sgd_float *vk = sgd_malloc(n_size);
	sgd_float *g  = sgd_malloc(n_size);

	for (int i = 0; i < n_size; i++)
	{
		mk[i] = vk[i] = 0.0;
	}

	sgd_float beta_1t = 1.0, beta_2t = 1.0;
	sgd_float alpha_k;

	sgd_float tmp_fx;
	for (int t = 0; t < para.iteration; t++)
	{
		beta_1t *= para.beta_1;
		beta_2t *= para.beta_2;

		alpha_k = para.alpha * sqrt(1.0 - beta_2t)/(1.0 - beta_1t);

		*fx = 0.0;
		for (int m = 0; m < m_size; m++)
		{
			tmp_fx = Evafp(instance, x, g, n_size, m);
			*fx += tmp_fx;

			for (int i = 0; i < n_size; i++)
			{
				mk[i] = para.beta_1*mk[i] + (1.0 - para.beta_1)*g[i];
				vk[i] = para.beta_2*vk[i] + (1.0 - para.beta_2)*(g[i] - mk[i])*(g[i] - mk[i]);

				x[i] = x[i] - alpha_k * mk[i]/(sqrt(vk[i]) + para.sigma);
				if (x[i] != x[i]) return SGD_NAN_VALUE;
			}
		}
		*fx /= m_size;

		if (Profp(instance, *fx, x, g, param, n_size, t)) return SGD_STOP;
		if (*fx < para.epsilon) return SGD_CONVERGENCE;
	}

	sgd_free(mk);
	sgd_free(vk);
	sgd_free(g);
	return SGD_REACHED_MAX_ITERATIONS;
}
